package com.analysis.utils;

import com.analysis.dto.Reference;
import com.analysis.dto.TencentChatResponse;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;

/**
 * 流式响应处理器
 * 用于处理腾讯云智能体的流式响应，提取最终答案
 */
public class StreamResponseProcessor {

    private static final Logger logger = LoggerFactory.getLogger(StreamResponseProcessor.class);

    private final ObjectMapper objectMapper;
    private final List<TencentChatResponse> responses = new ArrayList<>();
    private TencentChatResponse finalReply = null;
    private boolean isComplete = false;
    private int tokenCount = 0;
    // 参考来源信息收集
    private final List<Reference> references = new ArrayList<>();

    public StreamResponseProcessor(ObjectMapper objectMapper) {
        this.objectMapper = objectMapper;
    }

    /**
     * 处理流式响应数据
     *
     * @param eventType  事件类型
     * @param eventData  事件数据
     * @param onMessage  消息回调
     * @param onComplete 完成回调
     */
    public void processEvent(String eventType, String eventData,
                             Consumer<TencentChatResponse> onMessage,
                             Consumer<TencentChatResponse> onComplete) {
        try {
            // 修复可能的编码问题
            String fixedData = EncodingUtils.fixEncoding(eventData);

            TencentChatResponse response = objectMapper.readValue(fixedData, TencentChatResponse.class);
            response.setType(eventType);

            // 添加到响应列表
            responses.add(response);

            // 处理不同类型的事件
            switch (eventType) {
                case "reply":
                    handleReplyEvent(response, onMessage, onComplete);
                    break;
                case "token_stat":
                    handleTokenStatEvent(response, onMessage);
                    break;
                case "reference":
                    handleReferenceEvent(response, onMessage);
                    break;
                case "error":
                    handleErrorEvent(response, onMessage);
                    break;
                default:
                    logger.debug("未处理的事件类型: {}", eventType);
                    break;
            }

        } catch (Exception e) {
            logger.error("处理流式响应事件失败", e);
        }
    }

    /**
     * 处理回复事件
     */
    private void handleReplyEvent(TencentChatResponse response,
                                  Consumer<TencentChatResponse> onMessage,
                                  Consumer<TencentChatResponse> onComplete) {
        if (response.getPayload() != null && response.getPayload().getContent() != null) {
            // 检查是否为最终回复且不是用户自己的消息
            if (Boolean.TRUE.equals(response.getPayload().getIsFinal()) &&
                    Boolean.FALSE.equals(response.getPayload().getIsFromSelf())) {

                finalReply = response;
                isComplete = true;

                if (onComplete != null) {
                    onComplete.accept(response);
                }

                logger.info("收到最终AI回复: {}", response.getPayload().getContent());
            } else if (Boolean.TRUE.equals(response.getPayload().getIsFromSelf())) {
                logger.debug("收到用户消息: {}", response.getPayload().getContent());
            } else {
                logger.debug("收到中间回复: {}", response.getPayload().getContent());
            }

            if (onMessage != null) {
                onMessage.accept(response);
            }
        }
    }

    /**
     * 处理Token统计事件
     */
    private void handleTokenStatEvent(TencentChatResponse response, Consumer<TencentChatResponse> onMessage) {
        if (onMessage != null) {
            onMessage.accept(response);
        }
        if (response.getPayload() != null && response.getPayload().getTokenCount() != null && response.getPayload().getTokenCount() > 0) {
            tokenCount += response.getPayload().getTokenCount();
        }
    }

    /**
     * 处理参考来源事件
     */
    private void handleReferenceEvent(TencentChatResponse response, Consumer<TencentChatResponse> onMessage) {
        if (onMessage != null) {
            onMessage.accept(response);
        }
        // 处理参考来源信息
        if (response.getPayload() != null && response.getPayload().getReferences() != null) {
            references.addAll(response.getPayload().getReferences());
        } else {
            logger.debug("参考来源事件中没有payload信息");
        }
    }

    /**
     * 根据类型数字获取类型名称
     */
    private String getReferenceTypeName(Integer type) {
        if (type == null) {
            return "unknown";
        }

        switch (type) {
            case 1:
                return "knowledge"; // 知识库
            case 2:
                return "document"; // 文档
            case 3:
                return "web"; // 网页
            default:
                return "type_" + type;
        }
    }


    /**
     * 处理错误事件
     */
    private void handleErrorEvent(TencentChatResponse response, Consumer<TencentChatResponse> onMessage) {
        if (onMessage != null) {
            onMessage.accept(response);
        }

        logger.error("收到错误事件: {}", response.getError());
    }

    /**
     * 获取最终回复
     *
     * @return 最终回复，如果没有则返回null
     */
    public TencentChatResponse getFinalReply() {
        if (finalReply != null) {
            return finalReply;
        }

        // 如果没有找到is_final=true的回复，则取最后一个AI回复
        return responses.stream()
                .filter(response -> "reply".equals(response.getType()))
                .filter(response -> response.getPayload() != null)
                .filter(response -> Boolean.FALSE.equals(response.getPayload().getIsFromSelf()))
                .reduce((first, second) -> second)
                .orElse(null);
    }

    /**
     * 获取最终答案内容
     *
     * @return 最终答案内容，如果没有则返回null
     */
    public String getFinalAnswer() {
        TencentChatResponse reply = getFinalReply();
        if (reply != null && reply.getPayload() != null) {
            return reply.getPayload().getContent();
        }
        return null;
    }

    /**
     * 检查是否已完成
     *
     * @return 是否已完成
     */
    public boolean isComplete() {
        return isComplete;
    }

    /**
     * 获取所有响应
     *
     * @return 所有响应列表
     */
    public List<TencentChatResponse> getAllResponses() {
        return new ArrayList<>(responses);
    }

    /**
     * 获取回复响应
     *
     * @return 所有回复响应列表
     */
    public List<TencentChatResponse> getReplyResponses() {
        return responses.stream()
                .filter(response -> "reply".equals(response.getType()))
                .toList();
    }

    /**
     * 获取Token统计响应
     *
     * @return 所有Token统计响应列表
     */
    public List<TencentChatResponse> getTokenStatResponses() {
        return responses.stream()
                .filter(response -> "token_stat".equals(response.getType()))
                .toList();
    }

    /**
     * 获取收集到的参考来源信息
     *
     * @return 参考来源信息列表
     */
    public List<Reference> getReferences() {
        return new ArrayList<>(references);
    }

    /**
     * 重置处理器状态
     */
    public void reset() {
        responses.clear();
        finalReply = null;
        isComplete = false;
        tokenCount = 0;
        references.clear();
    }

    /**
     * 获取处理统计信息
     *
     * @return 统计信息字符串
     */
    public String getStatistics() {
        long replyCount = responses.stream().filter(r -> "reply".equals(r.getType())).count();
        long tokenStatCount = responses.stream().filter(r -> "token_stat".equals(r.getType())).count();
        long referenceCount = responses.stream().filter(r -> "reference".equals(r.getType())).count();
        long errorCount = responses.stream().filter(r -> "error".equals(r.getType())).count();

        return String.format("总事件数: %d, 回复事件: %d, Token统计事件: %d, 参考来源事件: %d, 错误事件: %d, 参考来源数: %d, 是否完成: %s",
                responses.size(), replyCount, tokenStatCount, referenceCount, errorCount, references.size(), isComplete);
    }

    public int getTokenCount() {
        return tokenCount;
    }

    public void setTokenCount(int tokenCount) {
        this.tokenCount = tokenCount;
    }
}
